import torch
import torch.nn as nn

from torch.nn import functional
# 两种特征提取网络 mobilenet 和 inception_resnet_v1
from .Mobilenet import MobileNetV1
from .Inception_ResNet_v1 import InceptionResnetV1


# 构造基于mobilenet特征提取网络的卷积网络
class mobilenet(nn.Module):
    def __init__(self, pretrained):            # pretrained为boolean类型 用以判断是否采用已经训练好的模型参数 这个根据题目要求来
        super(mobilenet, self).__init__()
        self.model = MobileNetV1()
        if pretrained:
            # 如果是采用预先训练好的mobilenet模型参数 那就加载
            state_dict = torch.load(r'../../model/backbone_weights_of_mobilenetv1.pth')
            self.model.load_state_dict(state_dict)

        # 删除分类所需要的卷积层
        del self.model.fc
        del self.model.avg

    def forward(self, x):
        x = self.model.stage1(x)
        x = self.model.stage2(x)
        x = self.model.stage3(x)
        return x

# 构造基于inception_resnet_v1特征提取网络的卷积网络
class inception_resnet(nn.Module):
    def __init__(self, pretrained):
        super(inception_resnet, self).__init__()
        self.model = InceptionResnetV1()
        if pretrained:
            # 如果是采用预先训练好的inception-resnet模型参数 那就加载
            state_dict = torch.load(r'../../model/backbone_weights_of_inception_resnetv1.pth')
            self.model.load_state_dict(state_dict)

    def forward(self, x):
        x = self.model.conv2d_1a(x)
        x = self.model.conv2d_2a(x)
        x = self.model.conv2d_2b(x)
        x = self.model.maxpool_3a(x)
        x = self.model.conv2d_3b(x)
        x = self.model.conv2d_4a(x)
        x = self.model.conv2d_4b(x)
        x = self.model.repeat_1(x)
        x = self.model.mixed_6a(x)
        x = self.model.repeat_2(x)
        x = self.model.mixed_7a(x)
        x = self.model.repeat_3(x)
        x = self.model.block8(x)
        return x

# 构建facenet类 用以实例化
class Facenet(nn.Module):
    def __init__(self, backbone="mobilenet", dropout_prob=0.5, embedding_size=128, pretrained=False):
        super(Facenet, self).__init__()
        # 主干特征提取网络采用哪一种
        if backbone == "mobilenet":
            self.backbone = mobilenet(pretrained)
            flat_shape = 1024
        elif backbone == "inception_resnetv1":
            self.backbone = inception_resnet(pretrained)
            flat_shape = 1792
        else:
            raise ValueError('Unsupported backbone - `{}`, Use mobilenet, inception_resnetv1.'.format(backbone))
        # 对特征层的后续操作
        self.avg = nn.AdaptiveAvgPool2d((1, 1))             # 论文中进行一个平均池化
        self.Dropout = nn.Dropout(dropout_prob)             # 防止过拟合 有0.5概率不被激活
        self.Connect = nn.Linear(flat_shape, embedding_size, bias=False)       # 最终是要输出128维的向量 以判断距离
        self.last_bn = nn.BatchNorm1d(embedding_size, eps=0.001, momentum=0.1, affine=True)         # 对128维向量进行归一化

    def forward(self, x):
        x = self.backbone(x)             # 通过主干特征提取网络 获取特征层
        x = self.avg(x)
        x = x.view(x.size(0), -1)        # resize    x.size(0)=batch size
        x = self.Dropout(x)
        x = self.Connect(x)
        x = self.last_bn(x)
        x = functional.normalize(x, p=2, dim=1)        # 标准化 除以2范数
        return x



